#!/usr/bin/env python3
"""
RNN翻译器训练脚本
英文到中文翻译器训练
"""

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.translator import RNTranslator

def create_sample_data():
    """创建示例训练数据"""
    # 英文句子
    english_sentences = [
        "Hello world",
        "How are you",
        "I love programming",
        "Machine learning is fun",
        "Deep learning models",
        "Natural language processing",
        "Artificial intelligence",
        "Computer vision",
        "Neural networks",
        "Data science",
        "Hello, how are you today?",
        "I am learning deep learning",
        "This is a translation system",
        "Python programming language",
        "PyTorch framework",
        "RNN models for translation",
        "Attention mechanism",
        "Sequence to sequence learning",
        "Encoder decoder architecture",
        "Language translation task"
    ]
    
    # 对应的中文翻译
    chinese_sentences = [
        "你好世界",
        "你好吗",
        "我喜欢编程",
        "机器学习很有趣",
        "深度学习模型",
        "自然语言处理",
        "人工智能",
        "计算机视觉",
        "神经网络",
        "数据科学",
        "你好，今天怎么样？",
        "我正在学习深度学习",
        "这是一个翻译系统",
        "Python编程语言",
        "PyTorch框架",
        "用于翻译的RNN模型",
        "注意力机制",
        "序列到序列学习",
        "编码器解码器架构",
        "语言翻译任务"
    ]
    
    return english_sentences, chinese_sentences

def main():
    """主训练函数"""
    print("=== RNN英文到中文翻译器训练 ===")
    
    # 创建翻译器实例
    translator = RNTranslator(
        max_length=30,
        hidden_size=128,  # 减小隐藏层大小以适应示例数据
        num_layers=2
    )
    
    # 获取训练数据
    en_sentences, zh_sentences = create_sample_data()
    
    print(f"训练数据量: {len(en_sentences)} 对句子")
    print("英文示例:", en_sentences[:3])
    print("中文示例:", zh_sentences[:3])
    
    # 训练模型（使用注意力机制）
    print("\n开始训练模型...")
    translator.train(
        en_sentences=en_sentences,
        zh_sentences=zh_sentences,
        epochs=50,  # 减少训练轮次
        batch_size=8,  # 减小批次大小
        learning_rate=0.001,
        use_attention=True
    )
    
    # 保存模型
    model_path = "models/translator_model.pth"
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    translator.save_model(model_path)
    print(f"\n模型已保存到: {model_path}")
    
    # 测试翻译
    print("\n=== 测试翻译 ===")
    test_sentences = [
        "Hello world",
        "How are you",
        "I love programming",
        "Machine learning"
    ]
    
    for en_sent in test_sentences:
        translation = translator.translate(en_sent)
        print(f"英文: {en_sent}")
        print(f"中文: {translation}")
        print("-" * 30)

if __name__ == "__main__":
    main()